;;; guile-openai --- An OpenAI API client for Guile
;;; Copyright © 2023 Andrew Whatson <whatson@tailcall.au>
;;;
;;; This file is part of guile-openai.
;;;
;;; guile-openai is free software: you can redistribute it and/or modify
;;; it under the terms of the GNU Affero General Public License as
;;; published by the Free Software Foundation, either version 3 of the
;;; License, or (at your option) any later version.
;;;
;;; guile-openai is distributed in the hope that it will be useful, but
;;; WITHOUT ANY WARRANTY; without even the implied warranty of
;;; MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
;;; Affero General Public License for more details.
;;;
;;; You should have received a copy of the GNU Affero General Public
;;; License along with guile-openai.  If not, see
;;; <https://www.gnu.org/licenses/>.

(define-module (openai image)
  #:use-module (openai api image)
  #:use-module (openai client)
  #:use-module (openai utils magick)
  #:use-module (ice-9 binary-ports)
  #:use-module (ice-9 match)
  #:use-module (srfi srfi-8)
  #:use-module (srfi srfi-9)
  #:use-module (srfi srfi-9 gnu)
  #:export (openai-default-image-size
            openai-default-image-format

            image?
            image-url
            image-file

            openai-image
            openai-image-edit
            openai-image-variation))

(define-once openai-default-image-size
  (make-parameter 512))

(define-once openai-default-image-format
  (make-parameter 'b64))

(define-record-type <Image>
  (%make-image url file)
  image?
  (url image-url)
  (file image-file))

(define (write-image-file img)
  (let* ((raw-data (fetch-image-data img))
         (png-rgba (magick-convert-image raw-data
                                         #:format "PNG"
                                         #:alpha-channel 'Activate)))
    (call-with-port (mkstemp "/tmp/guile-openai-XXXXXX")
      (lambda (port)
        (put-bytevector port png-rgba)
        (port-filename port)))))

(define (write-transparent-mask size)
  (receive (width height)
      (cond ((string=? size "256x256")
             (values 256 256))
            ((string=? size "512x512")
             (values 512 512))
            ((string=? size "1024x1024")))
    (let ((png-rgba (magick-create-image #:width width
                                         #:height height
                                         #:format "PNG"
                                         #:alpha-channel 'Activate
                                         #:background-color "white"
                                         #:background-alpha 0.0)))
      (call-with-port (mkstemp "/tmp/guile-openai-XXXXXX")
        (lambda (port)
          (put-bytevector port png-rgba)
          (port-filename port))))))

(define (make-image img)
  (let ((url (image-data-url img)))
    (%make-image (if (unspecified? url) #f url)
                 (write-image-file img))))

(define (print-image img port)
  (format port "#<Image: ~a>" (image-file img)))

(set-record-type-printer! <Image> print-image)

(define parse-image-size
  (match-lambda
    ((? unspecified? value) value)
    ((? string? value) value)
    ((or 256 '256x256) "256x256")
    ((or 512 '512x512) "512x512")
    ((or 1024 '1024x1024) "1024x1024")))

(define parse-image-format
  (match-lambda
    ((? unspecified? value) value)
    ((? string? value) value)
    ('url "url")
    ('b64 "b64_json")))

(define* (openai-image prompt #:key
                       (n      *unspecified*)
                       (size   (openai-default-image-size))
                       (format (openai-default-image-format))
                       (user   (openai-default-user)))
  "Send an image generation request.  Returns an image record.

The PROMPT must be a string.

The keyword arguments correspond to the request parameters described
in the image generation request documentation:

#:n - The number of images to generate, returned as multiple values.

#:size - A number specifying the generated image size, either 256,
512, or 1024.  The symbols `256x256', `512x512' and `1024x1024' are
also supported.

#:format - How the image data will be returned, either `url' (hosted
on a cloud server), or `b64' (embedded in the response message, the
default).  In either case, the image data will be fetched and stored
locally for display.

#:user - An optional username to associate with this request."
  (let* ((size (parse-image-size size))
         (format (parse-image-format format))
         (request (make-image-request prompt n size format user))
         (response (send-image-request request)))
    (apply values
           (map make-image (image-response-data response)))))

(define* (openai-image-edit image prompt #:key
                            (mask   *unspecified*)
                            (n      *unspecified*)
                            (size   (openai-default-image-size))
                            (format (openai-default-image-format))
                            (user   (openai-default-user)))
  "Send an image edit request.  Returns an image record.

The IMAGE can be an image record or a string path to the image file to
be varied.  The API requires the image to be a valid PNG file, less
than 4MB, and square.

The PROMPT must be a string describing the changes to be made to the
image.

The keyword arguments correspond to the request parameters described
in the image edit request documentation:

#:mask - An image record or a string path to an image file to use as
the edit mask.  The model will only edit pixels which are fully
transparent (ie. alpha 0) in the mask image.  The mask must have the
same dimensions as IMAGE.  The default is fully opaque, so no pixels
will be edited and a copy of the original image will be returned.

#:n - The number of images to generate, returned as multiple values.

#:size - A number specifying the generated image size, either 256,
512, or 1024.  The symbols `256x256', `512x512' and `1024x1024' are
also supported.

#:format - How the image data will be returned, either `url' (hosted
on a cloud server), or `b64' (embedded in the response message, the
default).  In either case, the image data will be fetched and stored
locally for display.

#:user - An optional username to associate with this request."
  (let* ((image (if (image? image) (image-file image) image))
         (mask (if (image? mask) (image-file mask) mask))
         (size (parse-image-size size))
         (format (parse-image-format format))
         (request (make-image-edit-request image mask prompt
                                           n size format user))
         (response (send-image-edit-request request)))
    (apply values
           (map make-image (image-response-data response)))))

(define* (openai-image-variation image #:key
                                 (n      *unspecified*)
                                 (size   (openai-default-image-size))
                                 (format (openai-default-image-format))
                                 (user   (openai-default-user)))
  "Send an image variation request.  Returns an image record.

The IMAGE can be an image record or a string path to the image file to
be varied.  The API requires the image to be a valid PNG file, less
than 4MB, and square.

The keyword arguments correspond to the request parameters described
in the image variation request documentation:

#:n - The number of images to generate, returned as multiple values.

#:size - A number specifying the generated image size, either 256,
512, or 1024.  The symbols `256x256', `512x512' and `1024x1024' are
also supported.

#:format - How the image data will be returned, either `url' (hosted
on a cloud server), or `b64' (embedded in the response message, the
default).  In either case, the image data will be fetched and stored
locally for display.

#:user - An optional username to associate with this request."
  (let* ((image (if (image? image) (image-file image) image))
         (size (parse-image-size size))
         (format (parse-image-format format))
         (request (make-image-variation-request image n size format user))
         (response (send-image-variation-request request)))
    (apply values
           (map make-image (image-response-data response)))))
